import nn
import numpy as np
import ops
import optimizor
import model
from datasets.MNIST.mnist.mnist import load_mnist


class DNN(nn.net):
    def __init__(self):
        super().__init__()
        self.layer1 = ops.MetaDense(input_channel=28 * 28, output_channel=512, init_lambda=0.8)
        self.layer2 = ops.ReLU()
        self.layer3 = ops.MetaDense(input_channel=512, output_channel=100, init_lambda=0.8)
        self.relu2 = ops.ReLU()
        self.batch = ops.BatchNormalization(gamma=512, beta=512)
        self.dense1 = ops.MetaDense(input_channel=100, output_channel=10, init_lambda=0.8)
        self.layer5 = ops.Softmax_cross_entropy()
        self.init_layers()

    def __call__(self, inputs, target):
        x = self.layer1(inputs)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.relu2(x)
        x = self.batch(x)
        x = self.dense1(x)
        outputs, loss_value = self.layer5(x, target)
        return outputs, loss_value

    def backward(self):
        grad = self.layer5.backward()
        gradx, gradw = self.dense1.backward(grad)
        grad = self.batch.backward(gradx)
        grad = self.relu2.backward(grad)
        grad_x2, grad_w2 = self.layer3.backward(grad)
        grad = self.layer2.backward(grad_x2)
        self.layer1.backward(grad)


if __name__ == "__main__":
    (train, train_lable), (test, test_lable) = load_mnist()
    train = train[:300]
    train_lable = train_lable[:300]
    test = test[:300]
    test_lable = test_lable[:300]
    train = list(map(lambda x: np.array([x]), train))
    true_y = np.zeros((len(train_lable), max(train_lable) + 1))
    for i in range(len(train_lable)):
        true_y[i][train_lable[i]] = 1
    train_lable = true_y
    test = list(map(lambda x: np.array([x]), test))
    true_y = np.zeros((len(test_lable), max(test_lable) + 1))
    for i in range(len(test_lable)):
        true_y[i][test_lable[i]] = 1
    test_lable = true_y

    mynet = DNN()
    model = model.Model(network=mynet)

    model.compile(
        optimizer=optimizor.AdaGrad(lr=0.01)
    )

    model.train(
        x=train,
        y=train_lable,
        epochs=20,
        batch_size=100
    )

    model.fit(
        x=test,
        y=test_lable
    )

    model.drawHistory()
